import xml.etree.cElementTree as ET

from utils.data_model import OneImage, RecBOX


def parse_xml(xml_file, ignore_list=None, replace_dct=None):
    if replace_dct is None:
        replace_dct = {}
    if ignore_list is None:
        ignore_list = []
    try:
        tree = ET.parse(xml_file)
        root = tree.getroot()
        width = int(root.findall('size')[0].find('width').text)
        height = int(root.findall('size')[0].find('height').text)
        one_image = OneImage(width, height)
        for obj in root.findall('object'):
            name = obj.find('name').text
            if ignore_list and name in ignore_list:
                continue
            if replace_dct and name in replace_dct.keys():
                name = replace_dct[name]
            bnd_box = obj.find('bndbox')
            x1 = int(bnd_box.find('xmin').text)
            y1 = int(bnd_box.find('ymin').text)
            x2 = int(bnd_box.find('xmax').text)
            y2 = int(bnd_box.find('ymax').text)
            box = RecBOX(name, x1, y1, x2, y2)
            one_image.add(box)
        return one_image
    except Exception as err:
        print(err)
        raise Exception('parse xml error')
